#!/usr/bin/env python3
import rospy
import threading
import torch
import func.visualization as v
from sensor_msgs.msg import Image
from uav_event.msg import CameraData, UInt32Mat
from cv_bridge import CvBridge
from rich import print


publisherInitialized = False
rawImgPub = None
egoMotionImgPub = None


subscriberInitialized = False
rawDataSub = None
egoMotionDataSub = None


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
br = CvBridge()


def rawDataSubCallback(msg: CameraData) -> None:
    global publisherInitialized, rawImgPub, device, br
    eventData = torch.tensor(msg.events, dtype = torch.long, device = device)
    if eventData.shape[0]:
        eventSize = len(msg.eventcol)
        eventData = eventData.reshape(eventData.shape[0] // eventSize, eventSize) # [sec, nsec, p, y, x]
        eventData = eventData.cpu().numpy()
        eventImg = v.white_bg(800, 1280)
        eventImg = v.add_events(eventImg, eventData[:, 2], eventData[:, 3], eventData[:, 4])
        if publisherInitialized:
            rawImgPub.publish(br.cv2_to_imgmsg(eventImg, encoding = "bgr8"))


def egoMotionSubCallback(msg: UInt32Mat) -> None:
    global publisherInitialized, egoMotionImgPub, device, br
    eventData = torch.tensor(msg.data, dtype = torch.long, device = device)
    if eventData.shape[0]:
        eventSize = len(msg.column)
        eventData = eventData.reshape(eventData.shape[0] // eventSize, eventSize) # [sec, nsec, p, y, x]
        eventData = eventData.cpu().numpy()
        eventImg = v.white_bg(80, 128)
        eventImg = v.add_events(eventImg, eventData[:, 2], eventData[:, 3], eventData[:, 4])
        if publisherInitialized:
            egoMotionImgPub.publish(br.cv2_to_imgmsg(eventImg, encoding = "bgr8"))


def subThreadFunc() -> None:
    global subscriberInitialized, rawDataSub, egoMotionDataSub
    rawDataSub = rospy.Subscriber("/uav_event/event_imu_data", CameraData, rawDataSubCallback)
    egoMotionDataSub = rospy.Subscriber("/uav_event/ego_motion_comp_res", UInt32Mat, egoMotionSubCallback)
    subscriberInitialized = True
    rospy.spin()


def pubThreadFunc() -> None:
    global publisherInitialized, rawImgPub, egoMotionImgPub
    rawImgPub = rospy.Publisher("/uav_event/raw_data_img", Image, queue_size = 1)
    egoMotionImgPub = rospy.Publisher("/uav_event/ego_motion_comp_img", Image, queue_size = 1)
    publisherInitialized = True


if __name__ == "__main__":
    rospy.init_node("view", anonymous = True)

    subT = threading.Thread(target = subThreadFunc)
    pubT = threading.Thread(target = pubThreadFunc)

    subT.start()
    pubT.start()

    subT.join()
    pubT.join()